package com.sparrow.common.lock;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 防止死锁的顺序锁
 * 线程A: lock(1) -> lock(2)、线程B: lock(2) -> lock(1)，在 tryLock() 失败后也都会退让 变为 1 -> 2 顺序，从而避免死锁
 * @param <K>
 */
public class LockChain<K extends Comparable> {
    private Map<K, ReentrantLock> lockMap = new ConcurrentHashMap<>();
    private final ThreadLocal<List<K>> threadHeldLocks = ThreadLocal.withInitial(ArrayList::new);

    public void lock(K id, Runnable task) {
        List<K> heldLocks = threadHeldLocks.get();
        ReentrantLock lock = lockMap.computeIfAbsent(id, k -> new ReentrantLock(false));
        if (lock.tryLock()) {
            heldLocks.add(id);
            try {
                task.run();
            } finally {
                if (lock.isHeldByCurrentThread()) {
                    lock.unlock();
                    heldLocks.remove(id);
                    if (heldLocks.isEmpty()) {
                        threadHeldLocks.remove();
                    }
                }
            }
        } else {//加锁失败可能遇到死锁 或者 没死锁但对方没释放锁
            //释放已持有的锁
            for (K e : heldLocks) {
                ReentrantLock l = lockMap.get(e);
                if (l.isHeldByCurrentThread()) {
                    l.unlock();
                }
            }
            // 构建新的锁顺序列表（去重 + 排序）
            List<K> allLocks = new ArrayList<>(heldLocks);
            if (!allLocks.contains(id)) {
                allLocks.add(id);
            }
            Collections.sort(allLocks);

            // 清空 heldLocks，重新构建
            heldLocks.clear();
            for (K e : allLocks) {
                ReentrantLock l = lockMap.get(e);
                l.lock(); // 阻塞加锁
                heldLocks.add(e); // 记录真正持有的锁
            }
            try {
                task.run();
            } finally {
                for (K e : heldLocks) {
                    ReentrantLock l = lockMap.get(e);
                    if (l.isHeldByCurrentThread()) {
                        l.unlock();
                    }
                }
                heldLocks.clear();
                threadHeldLocks.remove();
            }
        }
    }

}
